import asyncio
import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate


async def main():
    load_dotenv(override=True)

    prompt = ChatPromptTemplate.from_messages(
        [("system", "你叫做小智，是一个乐于助人的助手，请根据用户的输入的问题进行回答"),
         ("human", "{input}")],
    )
    model = init_chat_model(
        model="deepseek-chat",
        model_provider="deepseek",
        api_key=os.getenv("DEEPSEEK_API_KEY")

    )
    parser = StrOutputParser()
    chain = prompt | model | parser

    async for chunk in chain.astream({"input": "你好，我叫做陈明,请介绍一下南京"}):
        print(chunk, end="", flush=True)


if __name__ == '__main__':
    asyncio.run(main())
